# -*- coding: utf-8 -*-
import gc
import re
import sys
import numpy as np
import pandas as pd
import pyqtgraph as pg
from PyQt5.QtGui import QFont
from PyQt5.QtWidgets import QApplication
from pyqtgraph.exporters import ImageExporter
from pyqtgraph.icons import getGraphIcon
from PyQt5.QtCore import QPoint
import pyqtgraph.functions as fn
from pyqtgraph.graphicsItems.ScatterPlotItem import drawSymbol
from dataload.chartData import get_waferSummary
from src.sharedSrc import getSetting, rHEX, is_valid_hex
from PyQt5.QtCore import pyqtSignal

class WaferBlock(pg.ItemSample):
    """
    Used for changing square sizes in legends.
    """

    def paint(self, p, *args):
        opts = self.item.opts
        visible = self.item.isVisible()
        if not visible:
            icon = getGraphIcon("invisibleEye")
            p.drawPixmap(QPoint(1, 1), icon.pixmap(18, 18))
            return

        symbol = opts.get('symbol', None)
        if symbol is not None:
            p.translate(10, 15)
            drawSymbol(p, symbol, 20, fn.mkPen(opts['pen']), fn.mkBrush(opts['brush']))


class WaferMap(pg.GraphicsView):
    cdfSignal = pyqtSignal(str, str, str)
    def __init__(self):
        super().__init__()
        self.text_item = None
        self.measure = None
        self.plotlayout = pg.GraphicsLayout()
        self.setCentralWidget(self.plotlayout)
        self.view_list = []
        self.setBackground('white')
        self.settings = getSetting()


    def setWaferData(self, wafer_data: dict, file_name: str = '') -> None:
        """
        Establish the wafer data and visualize the wafer map.

        :param wafer_data: Data to be plotted
        :param file_name: Path to save the export file
        :return: None
        """
        gc.collect()  # Clear cache

        if len(wafer_data) == 0:
            return

        wafer_view = pg.ViewBox()
        self.pitem = pg.PlotItem(viewBox=wafer_view)
        wafer_precision = self.settings.wafer_precision
        # Used for setting the site squares
        legend = self.pitem.addLegend(offset=(40, 5), verSpacing=5, labelTextSize="15pt")
        # legend_value = self.pitem.addLegend(offset=(45, 5), verSpacing=5, labelTextSize="15pt")
        dataLists = []
        xy_data = wafer_data["Data"]
        sorted_keys = sorted(xy_data.keys())
        self.testName = wafer_data['TestName']
        self.scatter_items = []  # Store all scatter plot items

        for num in sorted_keys:
            xy_dict = xy_data[num]
            if is_valid_hex(num):
                num_split = num[0]
                color = self.settings.wafer_dict.get(num_split, rHEX())
            else:
                color = self.settings.wafer_dict.get(num, rHEX())
            test_name = xy_dict['name']
            # tip_func = f"XY: ({{x:.0f}}, {{y:.0f}}) \n {test_name}".format
            legend_string = f"{test_name}"
            self.measure = xy_dict['measure']
            self.spi = pg.ScatterPlotItem(
                symbol="s",
                pen=None,
                size=0.95,
                pxMode=False,
                hoverable=True,
                hoverPen=pg.mkPen('r', width=2),
                hoverSize=1,
            )
            self.spi.addPoints(x=xy_dict["x"], y=xy_dict["y"], brush=color, data=list(self.measure.values()))
            self.pitem.addItem(self.spi)

            self.scatter_items.append(self.spi)
            if wafer_data['Unit']:
                dataLists.extend(self.measure.values())
            legend.addItem(WaferBlock(self.spi), legend_string)

            text_items = [(x, y, self.measure[f"{x} {y}"]) for x, y in zip(xy_dict["x"], xy_dict["y"])]
            # Create a QFont object and set the desired font size
            font = QFont()
            font.setPointSize(self.settings.wfont_size)  # Set the font size to 12, for example
            if self.settings.parametric_map:
                for x, y, text in text_items:
                    text_item = pg.TextItem(text=text, anchor=(0.5, 0.5), color="#000000")
                    text_item.setPos(x, y)
                    text_item.setFont(font)
                    self.pitem.addItem(text_item)
        ratio = wafer_data["Info"]
        x_max, x_min, y_max, y_min = wafer_data["Bounds"]

        wafer_view.setLimits(xMin=x_min - 25, xMax=x_max + 25,
                             yMin=y_min - 25, yMax=y_max + 25,
                             maxXRange=(x_max - x_min + 100),
                             maxYRange=(y_max - y_min + 100),
                             minXRange=2, minYRange=2)
        wafer_view.setRange(xRange=(x_min - 5, x_max + 5),
                            yRange=(y_min - 5, y_max + 5),
                            disableAutoRange=False)
        wafer_view.setAspectLocked(lock=True, ratio=ratio)

        if self.settings.wafer_invert[self.settings.wafer_origin_index]['x']:
            wafer_view.invertX(True)
        if self.settings.wafer_invert[self.settings.wafer_origin_index]['y']:
            wafer_view.invertY(True)

        unit = wafer_data['Unit']
        if unit:
            dataLists = list(map(float, dataLists))
            max_value, min_value, mean_value = np.max(dataLists), np.min(dataLists), np.round(np.mean(dataLists),
                                                                                              wafer_precision.get(unit,
                                                                                                                  3))
            self.plotlayout.addLabel(
                f"{wafer_data['TestName']} {wafer_data['Signal']} Max: {max_value} Mean: {mean_value} Min: {min_value}",
                row=0, col=0,
                rowspan=1, colspan=2, size="16pt", color="#000000")
        else:
            self.plotlayout.addLabel(
                f"{wafer_data['TestName']} {wafer_data['Signal']}",
                row=0, col=0,
                rowspan=1, colspan=2, size="16pt", color="#000000")
        self.pitem.getAxis("left").setLabel(f"Test Value ({unit})" if unit else "Test Value")
        self.plotlayout.addItem(self.pitem, row=1, col=0, rowspan=1, colspan=2)
        self.view_list.append(wafer_view)
        # self.spi.sigClicked.connect(self.on_click)  # 点击

        if file_name != '':
            self.save_plot(file_name)

        # Add a click event to the ScatterPlotItem
        for spi in self.scatter_items:
            spi.sigClicked.connect(self.onPointClicked)

    def save_plot(self, filename: str) -> None:
        self.resize(2100, 1200)  # Example size
        self.plotlayout.resize(2100, 1200)
        exporter = ImageExporter(self.plotlayout.scene())
        exporter.export(filename)

    def onPointClicked(self, plot, points):
        """
        Handle the event when a point on the scatter plot is clicked.

        :param plot: The ScatterPlotItem that was clicked
        :param points: The points that were clicked
        """
        if points and self.testName:
            # Get the first clicked point
            point = points[0]
            pos = point.pos()
            x, y = int(pos.x()), int(pos.y())
            self.cdfSignal.emit(self.testName, str(x), str(y))


    def leaveEvent(self, event) -> None:
        if self.text_item:
            self.pitem.removeItem(self.text_item)
            self.text_item = None
        super().leaveEvent(event)


def test_measure(test_name):
    df = pd.read_csv(r"C:\Users\29762\Downloads\formingyield11_20240905-111131_datalog_20240905115746.csv", dtype=str,
                     low_memory=False)

    # Extract columns containing "Measure"
    measure_cols = [col for col in df.columns if 'Measure' in col]

    if not measure_cols:
        return {}

    # Filter rows by test name
    filtered_df = df[df['TestName'] == test_name]
    if filtered_df.empty:
        return {}

    # Unit conversion dictionary
    unit_conversion = {
        "mV": 1, "V": 1000, "uV": 0.001, "nV": 0.000001,
        "mA": 1, "A": 1000, "uA": 0.001, "nA": 0.000001,
    }

    # Extract and convert values
    all_measure = []
    for measure in measure_cols:
        values_with_units = filtered_df[measure].dropna()
        for value in values_with_units:
            value_and_unit = extract_value_and_unit(value)
            if value_and_unit:
                unit = value_and_unit['unit']
                value = float(value_and_unit['value'])
                converted_value = value * unit_conversion.get(unit, 1)
                all_measure.append(converted_value)

    if not all_measure:
        return {}
    # print(len(all_measure))

    measure_ndarray = np.asarray(all_measure)
    sorted_data = np.sort(measure_ndarray)
    # length = len(all_measure)
    # mean_value = np.mean(measure_ndarray)
    cdf = np.arange(1, len(sorted_data) + 1) / len(sorted_data)
    RD = {}
    RD_Ticks = {}
    for i in range(len(cdf)):
        RD[i] = cdf[i]
        RD_Ticks[i] = (i, str(cdf[i]))
    # print(mean_value)
    # # Calculate min, max, and binning
    min_value = min(all_measure)
    max_value = max(all_measure)

    # print(min_value, ' ', max_value)
    # logspace_values = np.logspace(min_value, max_value, num=5, base=10)
    # print(logspace_values)
    # print(min_value, ' ', max_value)
    # lower_length = len(np.where(measure_ndarray < mean_value)[0])
    # upper_length = length - lower_length
    # print(lower_length, ' ', upper_length)

    # space_value = int(max_value - min_value) + 1
    # space = 1 if space_value < 100 else space_value // 100
    #
    # Relative Distribution (RD)
    # RD = {i: sum(1 for value in all_measure if value < (space * i + min_value))
    #       for i in range(space_value)}
    # RD_Ticks = {i: (i, str(space * i + min_value)) for i in range(space_value)}

    # Normalized Distribution (ND)
    split = 5
    split_edges = np.linspace(min_value, max_value, split + 1)
    ND = {}
    ND_Ticks = {}

    for i in range(split):
        lower_bound = split_edges[i]
        upper_bound = split_edges[i + 1]
        set_color = f'split_{i}'
        measure_cnt = sum(1 for value in all_measure
                          if lower_bound <= value < upper_bound or (i == split - 1 and value == upper_bound))
        ND[set_color] = measure_cnt
        ND_Ticks[set_color] = (set_color, str(measure_cnt))

    # Compile results
    bin_data = {
        "ND": {1: ND},
        "RD": {1: RD},
        "ND_Ticks": ND_Ticks,
        "RD_Ticks": RD_Ticks,
        "min": min_value,
        "max": max_value,
    }
    return bin_data


def get_xyMeasure(test_name, x, y):
    df = pd.read_csv(r"D:\dataset\new3\auto_20240829-142329_datalog_19700101080000.csv", dtype=str, low_memory=False)
    measure_list = [col for col in df.columns if 'Measure' in col]
    XCoord = df[df['Force'] == 'XCoord']
    YCoord = df[df['Force'] == 'YCoord']
    measure = ''
    unit_conversion = {
        "mV": 1,
        "V": 1000,
        "uV": 0.001,
        "nV": 0.000001,
        "mA": 1,
        "A": 1000,
        "uA": 0.001,
        "nA": 0.000001,
    }
    for col in measure_list:
        if x == XCoord[col].values[-1] and y == YCoord[col].values[-1]:
            measure = col
    if measure:
        measure_df = df[measure]
        testNames = df['TestName']
        all_measure = []
        for name, mea in zip(testNames, measure_df):
            if name == test_name:
                value_and_unit = extract_value_and_unit(mea)
                if value_and_unit:
                    u = value_and_unit['unit']
                    v = value_and_unit['value']
                    conversion = unit_conversion.get(u, 1)
                    parsed_value = float(v) * conversion
                    all_measure.append(parsed_value)
        if len(all_measure) == 0:
            return {}
        min_value, max_value = min(all_measure), max(all_measure)
        space_value = int(max_value - min_value) + 1
        space = 1 if space_value < 100 else space_value // 100
        RD = {}
        RD_Ticks = {}
        for i in range(space_value):
            bin_min = space * i + min_value
            RD[i] = sum(1 for j in all_measure if j < bin_min)
            RD_Ticks[i] = (i, str(bin_min))
        bin_data = {
            "RD": {1: RD},
            "RD_Ticks": RD_Ticks,
            "min": min_value,
            "max": max_value,
        }
        return bin_data
    else:
        return {}

def extract_value_and_unit(input_string):
    # 正则表达式匹配必须有一个数值和一个单位
    regex = r'^([+-]?\d*\.?\d+)\s*(mV|mA|uA|nA|A|uV|nV|V)$'
    match = re.match(regex, input_string)

    if match:
        # 第一个括号匹配的是数值部分
        value = float(match.group(1))
        # 第二个括号匹配的是单位部分
        unit = match.group(2)

        return {"value": value, "unit": unit}
    else:
        # 如果输入格式不匹配，则返回None或其他适当的错误处理
        return None


def mapSplitData(df: pd, testName: str, signal: str, precision: dict, split: int = 5) -> dict:
    """
    get bin from csv
    :param split:
    :param precision:
    :param testName:
    :param signal:
    :param df: pandas read csv, type Dataframe
    :return: process bin data, type dict
    """

    # Filter DataFrame once instead of multiple times
    data = df[(df['TestName'] == testName) & (df['Signal'] == signal)]
    data = data.filter(regex='^Measure')
    if data.empty:
        return {}
    # units = ["mV", "mA", "uA", "nA", "A", "uV", "nV", "V"]
    # unit_conversion = [1, 1, 10e-3, 10e-6, 1e3, 10e-3, 10e-6, 1e3]
    unit_conversion = {
        "mV": 1,
        "V": 1000,
        "uV": 0.001,
        "nV": 0.000001,
        "mA": 1,
        "A": 1000,
        "uA": 0.001,
        "nA": 0.000001,
    }
    set_color = 'value'
    x_coord, y_coord = [], []
    x_coord_df, y_coord_df = df[df['Force'] == 'XCoord'], df[df['Force'] == 'YCoord']
    xy_measure, mvaMeasure = {}, {}
    x_mva, y_mva = [], []
    unit = None
    for v in data:
        value = data[v].values[0]
        if pd.isna(value) or not value:
            continue
        value_and_unit = extract_value_and_unit(value)
        x, y = x_coord_df[v].values, y_coord_df[v].values
        key = f"{x[0]} {y[0]}"
        if value_and_unit:
            u = value_and_unit['unit']
            v = value_and_unit['value']

            conversion = unit_conversion.get(u, 1)
            parsed_value = float(v) * conversion
            unit = 'mA' if "A" in u else 'mV'
            x_mva.extend(x)
            y_mva.extend(y)
            mvaMeasure[key] = parsed_value
        else:
            # 非电压 电流组
            x_coord.extend(x)
            y_coord.extend(y)
            xy_measure[key] = str(value)
            set_color = value

    if xy_measure:
        x_coord, y_coord = list(map(int, x_coord)), list(map(int, y_coord))
        x_min, x_max, y_min, y_max = min(x_coord), max(x_coord), min(y_coord), max(y_coord)
    elif mvaMeasure:
        # 转换坐标
        x_mva, y_mva = list(map(int, x_mva)), list(map(int, y_mva))
        # 查找x列表和y列表中的最大值与最小值
        x_min, x_max, y_min, y_max = min(x_mva), max(x_mva), min(y_mva), max(y_mva)
    else:
        return {}

    sample_wafer_data = {
        "Data": {

        },
        "Info": 1.0,
        "Unit": unit,
        "Bounds": (x_max, x_min, y_max, y_min),
        "TestName": testName,
        "Signal": signal
    }
    if xy_measure:
        sample_wafer_data["Data"][set_color] = {"x": x_coord, "y": y_coord, "measure": xy_measure,
                                                'name': f'{testName} {signal}'}
    else:
        precision = precision.get(unit, 3)
        list_measure = np.array(list(mvaMeasure.values()))
        min_value, max_value = list_measure.min(), list_measure.max()
        split_edges = np.linspace(min_value, max_value, split + 1)

        for i in range(split):
            xc, yc, measure = [], [], {}
            set_color = f'split_{i}'
            lower_bound = split_edges[i]
            upper_bound = split_edges[i + 1]

            # Assign each value to the appropriate bin
            for key, value in mvaMeasure.items():
                x, y = map(int, key.split())
                if lower_bound <= value < upper_bound:
                    xc.append(x)
                    yc.append(y)
                    measure[key] = str(round(value, precision))
                elif i == split - 1 and value == upper_bound:
                    # Ensure the maximum value is included in the last bin
                    xc.append(x)
                    yc.append(y)
                    measure[key] = str(round(value, precision))

            if measure:  # Only add non-empty splits
                sample_wafer_data["Data"][set_color] = {
                    "x": xc,
                    "y": yc,
                    "measure": measure,
                    'name': f'{set_color}: [{lower_bound:.{precision}f} ~ {upper_bound:.{precision}f}]'
                }

    return sample_wafer_data


def main():
    df = pd.read_csv(r"C:\Users\29762\Downloads\formingyield11_20240905-111131_datalog_20240905115746.csv", dtype=str,
                     low_memory=False)

    test_name = 'Icell0'
    signal = '00_HR'
    sample_wafer_data = mapSplitData(df, test_name, signal, {1: 1})

    app = QApplication(sys.argv)
    wafer_map = WaferMap()
    wafer_map.setWaferData(sample_wafer_data, file_name=r"D:\dataset\new3\export\tetimg\OS_i_clk_map.jpg")
    wafer_map.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()
    # get_xyMeasure(1, 1, 1)
    test_measure('Icell0')
